
segrap_subset = {
    'Brain': [1, 2, 3, 4, 5, 6, 7, 8, 9],
    "BrainStem": 2,
    "Chiasm": 3,
    "TemporalLobe_L": [4, 6],
    "TemporalLobe_R": [5, 7],
    "Hippocampus_L": [8, 6],
    "Hippocampus_R": [9, 7],
    'Eye_L': [10, 12],
    'Eye_R': [11, 13],
    "Lens_L": 12,
    "Lens_R": 13,
    "OpticNerve_L": 14,
    "OpticNerve_R": 15,
    "MiddleEar_L": [18, 16, 20, 24, 28, 30],
    "MiddleEar_R": [19, 17, 21, 25, 29, 31],
    "IAC_L": 18,
    "IAC_R": 19,
    "TympanicCavity_L": [22, 20],
    "TympanicCavity_R": [23, 21],
    "VestibulSemi_L": [26, 24],
    "VestibulSemi_R": [27, 25],
    "Cochlea_L": 28,
    "Cochlea_R": 29,
    "ETbone_L": [32, 30],
    "ETbone_R": [33, 31],
    "Pituitary": 34,
    "OralCavity": 35,
    "Mandible_L": 36,
    "Mandible_R": 37,
    "Submandibular_L": 38,
    "Submandibular_R": 39,
    "Parotid_L": 40,
    "Parotid_R": 41,
    "Mastoid_L": 42,
    "Mastoid_R": 43,
    "TMjoint_L": 44,
    "TMjoint_R": 45,
    "SpinalCord": 46,
    "Esophagus": 47,
    "Larynx": [48, 49, 50, 51],
    "Larynx_Glottic": 49,
    "Larynx_Supraglot": 50,
    "PharynxConst": [51, 52],
    "Thyroid": 53,
    "Trachea": 54}


from light_training.preprocessing.preprocessors.preprocessor_multiinput_and_region_01norm_first import MultiInputAndRegionPreprocessor 
import numpy as np 
import pickle 
import json 


def process_train():
    base_dir = "./data/raw_data/"
    image_dir = "SegRap2023_Training_Set_120cases"
    data_filenames = ["image.nii.gz", "image_contrast.nii.gz"]

    seg_filename = "seg.nii.gz"
    preprocessor = MultiInputAndRegionPreprocessor(base_dir=base_dir, 
                                    image_dir=image_dir,
                                    data_filenames=data_filenames,
                                    seg_filename=seg_filename,
                                    norm_clip_min=-175,
                                    norm_clip_max=250
                                   )
    
    # out_spacing = [3.0, 0.54199219, 0.54199219]
    out_spacing = [3.0, 1.0, 1.0]
    output_dir = "./data/lowres/train/"

    preprocessor.run(output_spacing=out_spacing, 
                     output_dir=output_dir, 
                     all_labels_dict=segrap_subset,
                     num_processes=32,
                    )

def process_val():
    # fullres spacing is [0.5        0.70410156 0.70410156]
    # median_shape is [602.5 516.5 516.5]
    base_dir = "./data/raw_data/Val"
    image_dir = "img"
    preprocessor = DefaultPreprocessor(base_dir=base_dir, 
                                    image_dir=image_dir,
                                    label_dir=None,
                                   )

    out_spacing = [0.5, 0.70410156, 0.70410156]

    with open("./data_analysis_result.txt", "r") as f:
        content = f.read().strip("\n")
        print(content)
    content = eval(content)
    foreground_intensity_properties_per_channel = content["intensity_statistics_per_channel"]

    output_dir = "./data/fullres/val_test/"
    preprocessor.run(output_spacing=out_spacing, 
                     output_dir=output_dir,
                     all_labels=[1, ],
                     foreground_intensity_properties_per_channel=foreground_intensity_properties_per_channel,
                     num_processes=16)

def process_val_semi():
    # fullres spacing is [0.5        0.70410156 0.70410156]
    # median_shape is [602.5 516.5 516.5]
    base_dir = "./data/raw_data/Val_semi_postprocess"
    image_dir = "img"
    preprocessor = DefaultPreprocessor(base_dir=base_dir, 
                                    image_dir=image_dir,
                                    label_dir="gt",
                                   )

    out_spacing = [0.5, 0.70410156, 0.70410156]

    with open("./data_analysis_result.txt", "r") as f:
        content = f.read().strip("\n")
        print(content)
    content = eval(content)
    foreground_intensity_properties_per_channel = content["intensity_statistics_per_channel"]

    output_dir = "./data/fullres/val_semi_postprocess/"
    preprocessor.run(output_spacing=out_spacing, 
                     output_dir=output_dir,
                     all_labels=[1, ],
                     foreground_intensity_properties_per_channel=foreground_intensity_properties_per_channel)


def plan():
    base_dir = "./data/raw_data/"
    image_dir = "SegRap2023_Training_Set_120cases"
    data_filenames = ["image.nii.gz", "image_contrast.nii.gz"]

    seg_filename = "seg.nii.gz"
    preprocessor = MultiInputAndRegionPreprocessor(base_dir=base_dir, 
                                    image_dir=image_dir,
                                    data_filenames=data_filenames,
                                    seg_filename=seg_filename
                                   )

    preprocessor.run_plan()

if __name__ == "__main__":

    # plan()

    process_train()
    # import time 
    # s = time.time()
    # process_val()
    # e = time.time()

    # print(f"preprocessing time is {e - s}")

 